#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UCEIS模型诊断可视化工具
提供混淆矩阵、训练曲线等可视化功能
"""

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from sklearn.metrics import confusion_matrix, classification_report
from pathlib import Path
import json
from datetime import datetime
import torch

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class UCEISVisualizer:
    """UCEIS模型可视化器"""

    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # 设置绘图风格
        plt.style.use('seaborn-v0_8')
        sns.set_palette("husl")

    def plot_confusion_matrix(self, y_true, y_pred, epoch=None, split='val',
                            class_names=None, normalize=False):
        """
        绘制混淆矩阵

        Args:
            y_true: 真实标签
            y_pred: 预测标签
            epoch: 当前epoch
            split: 数据集分割 ('train', 'val', 'test')
            class_names: 类别名称
            normalize: 是否归一化
        """
        if class_names is None:
            # 8分类UCEIS评分
            class_names = [f'{i+1}分' for i in range(8)]

        # 计算混淆矩阵
        cm = confusion_matrix(y_true, y_pred)

        if normalize:
            cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
            title = f'归一化混淆矩阵 - {split.upper()}'
            fmt = '.2f'
        else:
            title = f'混淆矩阵 - {split.upper()}'
            fmt = 'd'

        # 创建图形
        plt.figure(figsize=(10, 8))

        # 绘制热力图
        sns.heatmap(cm, annot=True, fmt=fmt, cmap='Blues',
                   xticklabels=class_names, yticklabels=class_names,
                   cbar_kws={'label': '样本数量'})

        plt.title(f'{title} (Epoch {epoch})' if epoch else title,
                 fontsize=16, fontweight='bold')
        plt.xlabel('预测标签', fontsize=12)
        plt.ylabel('真实标签', fontsize=12)
        plt.xticks(rotation=45)
        plt.yticks(rotation=0)
        plt.tight_layout()

        # 保存图像
        if epoch:
            filename = f'confusion_matrix_{split}_epoch_{epoch}.png'
        else:
            filename = f'confusion_matrix_{split}_final.png'

        filepath = self.output_dir / filename
        plt.savefig(filepath, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"混淆矩阵已保存: {filepath}")

        return cm

    def plot_training_curves(self, train_history, val_history):
        """
        绘制训练曲线

        Args:
            train_history: 训练历史字典
            val_history: 验证历史字典
        """
        epochs = range(1, len(train_history['loss']) + 1)

        # 创建2x2子图
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle('UCEIS模型训练曲线', fontsize=16, fontweight='bold')

        # 损失曲线
        axes[0, 0].plot(epochs, train_history['loss'], 'b-', label='训练损失', linewidth=2)
        axes[0, 0].plot(epochs, val_history['loss'], 'r-', label='验证损失', linewidth=2)
        axes[0, 0].set_title('损失曲线')
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Loss')
        axes[0, 0].legend()
        axes[0, 0].grid(True, alpha=0.3)

        # 准确率曲线
        axes[0, 1].plot(epochs, train_history['accuracy'], 'b-', label='训练准确率', linewidth=2)
        axes[0, 1].plot(epochs, val_history['accuracy'], 'r-', label='验证准确率', linewidth=2)
        axes[0, 1].set_title('准确率曲线')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Accuracy')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)

        # 平衡准确率曲线
        axes[1, 0].plot(epochs, train_history['balanced_accuracy'], 'b-',
                       label='训练平衡准确率', linewidth=2)
        axes[1, 0].plot(epochs, val_history['balanced_accuracy'], 'r-',
                       label='验证平衡准确率', linewidth=2)
        axes[1, 0].set_title('平衡准确率曲线')
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('Balanced Accuracy')
        axes[1, 0].legend()
        axes[1, 0].grid(True, alpha=0.3)

        # F1-macro曲线
        axes[1, 1].plot(epochs, train_history['f1_macro'], 'b-',
                       label='训练F1-macro', linewidth=2)
        axes[1, 1].plot(epochs, val_history['f1_macro'], 'r-',
                       label='验证F1-macro', linewidth=2)
        axes[1, 1].set_title('F1-macro曲线')
        axes[1, 1].set_xlabel('Epoch')
        axes[1, 1].set_ylabel('F1-macro')
        axes[1, 1].legend()
        axes[1, 1].grid(True, alpha=0.3)

        plt.tight_layout()

        # 保存图像
        filepath = self.output_dir / 'training_curves.png'
        plt.savefig(filepath, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"训练曲线已保存: {filepath}")

    def plot_class_distribution(self, class_counts, class_names=None, title="类别分布"):
        """
        绘制类别分布图

        Args:
            class_counts: 类别计数列表
            class_names: 类别名称
            title: 图表标题
        """
        if class_names is None:
            class_names = [f'{i+1}分' for i in range(len(class_counts))]

        plt.figure(figsize=(12, 6))

        # 创建柱状图
        bars = plt.bar(class_names, class_counts, alpha=0.7,
                      color=sns.color_palette("husl", len(class_counts)))

        # 在柱子上添加数值标签
        for bar, count in zip(bars, class_counts):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(class_counts)*0.01,
                    f'{count}', ha='center', va='bottom', fontweight='bold')

        plt.title(title, fontsize=16, fontweight='bold')
        plt.xlabel('UCEIS评分', fontsize=12)
        plt.ylabel('样本数量', fontsize=12)
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3, axis='y')
        plt.tight_layout()

        # 保存图像
        filepath = self.output_dir / 'class_distribution.png'
        plt.savefig(filepath, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"类别分布图已保存: {filepath}")

    def plot_3class_mapping_performance(self, labels_true, labels_pred, epoch=None, split='val'):
        """
        绘制3分类映射性能图

        Args:
            labels_true: 8分类真实标签
            labels_pred: 8分类预测标签
            epoch: 当前epoch
            split: 数据集分割
        """
        def map_8_to_3_class(score_8_class_idx):
            """将8分类映射到3分类"""
            if score_8_class_idx <= 2:  # 0-2 -> 轻度 (1-3分)
                return 0
            elif score_8_class_idx <= 5:  # 3-5 -> 中度 (4-6分)
                return 1
            else:  # 6-7 -> 重度 (7-8分)
                return 2

        # 转换为3分类
        labels_true_3class = [map_8_to_3_class(label) for label in labels_true]
        labels_pred_3class = [map_8_to_3_class(label) for label in labels_pred]

        class_names = ['轻度(1-3分)', '中度(4-6分)', '重度(7-8分)']

        # 计算混淆矩阵
        cm = confusion_matrix(labels_true_3class, labels_pred_3class)

        # 归一化
        cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

        # 创建图形
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

        # 原始计数混淆矩阵
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                   xticklabels=class_names, yticklabels=class_names,
                   cbar_kws={'label': '样本数量'}, ax=ax1)
        ax1.set_title(f'3分类混淆矩阵 - {split.upper()} (Epoch {epoch})' if epoch
                     else f'3分类混淆矩阵 - {split.upper()}', fontweight='bold')
        ax1.set_xlabel('预测标签')
        ax1.set_ylabel('真实标签')

        # 归一化混淆矩阵
        sns.heatmap(cm_normalized, annot=True, fmt='.2f', cmap='Blues',
                   xticklabels=class_names, yticklabels=class_names,
                   cbar_kws={'label': '归一化比例'}, ax=ax2)
        ax2.set_title(f'归一化3分类混淆矩阵 - {split.upper()} (Epoch {epoch})' if epoch
                     else f'归一化3分类混淆矩阵 - {split.upper()}', fontweight='bold')
        ax2.set_xlabel('预测标签')
        ax2.set_ylabel('真实标签')

        plt.tight_layout()

        # 保存图像
        if epoch:
            filename = f'3class_confusion_matrix_{split}_epoch_{epoch}.png'
        else:
            filename = f'3class_confusion_matrix_{split}_final.png'

        filepath = self.output_dir / filename
        plt.savefig(filepath, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"3分类混淆矩阵已保存: {filepath}")

        return cm, cm_normalized

    def save_diagnostics_report(self, report_data):
        """
        保存诊断报告

        Args:
            report_data: 报告数据字典
        """
        # 添加时间戳
        report_data['timestamp'] = datetime.now().isoformat()

        # 保存报告
        filepath = self.output_dir / 'diagnostics_report.json'
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(report_data, f, indent=2, ensure_ascii=False)

        print(f"诊断报告已保存: {filepath}")

    def create_epoch_summary(self, epoch, train_metrics, val_metrics,
                           y_true_train, y_pred_train, y_true_val, y_pred_val):
        """
        创建每个epoch的总结可视化

        Args:
            epoch: 当前epoch
            train_metrics: 训练指标
            val_metrics: 验证指标
            y_true_train, y_pred_train: 训练集真实和预测标签
            y_true_val, y_pred_val: 验证集真实和预测标签
        """
        # 生成混淆矩阵
        self.plot_confusion_matrix(y_true_train, y_pred_train, epoch, 'train')
        self.plot_confusion_matrix(y_true_val, y_pred_val, epoch, 'val')

        # 生成3分类混淆矩阵
        self.plot_3class_mapping_performance(y_true_train, y_pred_train, epoch, 'train')
        self.plot_3class_mapping_performance(y_true_val, y_pred_val, epoch, 'val')

        # 保存当前epoch的指标
        epoch_report = {
            'epoch': epoch,
            'train_metrics': train_metrics,
            'val_metrics': val_metrics,
            'timestamp': datetime.now().isoformat()
        }

        filepath = self.output_dir / f'epoch_{epoch}_report.json'
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(epoch_report, f, indent=2, ensure_ascii=False)